import json
import os
import re
from typing import Any, List


def extract_code(input: str) -> str | None:
    # 使用正则表达式查找Java代码块
    pattern = r"```(java)?(.*?)```"
    matches = re.findall(pattern, input, re.DOTALL)
    for match in matches:
        if "@Test" in match[1]:
            return match[1].strip()
    return None


def extract_json(input: str, expected_schema=None) -> list[Any] | None | Any:
    # 用于存储所有有效的 JSON 对象
    json_objects = []

    # 1. 如果整个字符串是 JSON 格式，则直接解析
    try:
        # 尝试直接解析整个字符串
        parsed_json = json.loads(input)
        if expected_schema and not validate_json_structure(parsed_json, expected_schema):
            return None, True
        json_objects.append(parsed_json)
    except json.JSONDecodeError:
        pass  # 如果整个字符串无法解析为 JSON，就继续尝试其他方式

    # 2. 提取包裹在三重反引号中的部分，并尝试解析每个部分
    # 正则表达式匹配 ` ```json ... ``` ` 或者 ` ``` ... ``` `
    code_block_pattern = r'```(json)?(.*?)```'  # 匹配带 json 或没有标识的三重反引号内容
    code_blocks = re.findall(code_block_pattern, input, flags=re.DOTALL)  # `re.DOTALL` 允许匹配换行符
    find_json = False

    for block in code_blocks:
        block_content = block[1].strip()  # 获取反引号中的内容并去除前后空白
        try:
            # 尝试解析每个代码块中的内容
            parsed_json = json.loads(block_content)
            find_json = True
            if expected_schema and not validate_json_structure(parsed_json, expected_schema):
                continue  # 如果格式或类型不符合预期，跳过此 JSON
            json_objects.append(parsed_json)
        except json.JSONDecodeError:
            pass  # 如果无法解析为 JSON，就忽略这个代码块

    # 如果有有效的 JSON 数据，返回 JSON 对象列表，否则返回 None
    return json_objects if json_objects else None, find_json


def validate_json_structure(json_obj, expected_schema):
    """
    验证 JSON 对象的结构和类型是否符合预期。

    :param json_obj: 要验证的 JSON 对象（字典）。
    :param expected_schema: 预期的 JSON 结构，格式为 {key: expected_type}，例如：
                            {"name": str, "value": int, "data": list}
    :return: 如果所有键的类型符合预期，返回 True；否则返回 False。
    """
    # 遍历预期结构中的每一个键
    if json_obj is None:
        return False
    for key, expected_type in expected_schema.items():
        # 检查键是否存在且类型匹配
        if key not in json_obj:
            print(f"Missing key: {key}")
            return False

        if not isinstance(json_obj[key], expected_type):
            print(f"Incorrect type for key '{key}': Expected {expected_type}, got {type(json_obj[key])}")
            return False

    return True

def add_record(log_dir, name, ret, find_bugs, coverage, mutation_score, line_coverage_rates, branch_coverage_rates,
               mutation_scores, total_lines=set(), covered_lines=set(), mutation_info={}, update=False):
    with open(os.path.join(log_dir, "record.jsonl"), "a") as f:
        f.write(json.dumps(
            {"update": update, "name": name, "ret": ret, "find_bugs": find_bugs, "coverage": coverage, "mutation_score": mutation_score,
             "line_coverage_rates": line_coverage_rates, "branch_coverage_rates": branch_coverage_rates,
             "mutation_scores": mutation_scores, "total_lines": list(total_lines), "covered_lines": list(covered_lines),
             "mutation_info": mutation_info}) + "\n")


def add_test_case(log_dir, name, find_bugs, test_point, test_case, signature, mutants_list=[]):
    with open(os.path.join(log_dir, "test_case.jsonl"), "a") as f:
        f.write(json.dumps({"signature": signature, "name": name, "find_bugs": find_bugs, "test_point": test_point,
                             "test_case": test_case, "mutants_list": mutants_list}) + "\n")

if __name__ == '__main__':
    # 测试 extract_json 函数
    input_text = """
    ```json
    {
        "name": "Alice",
        "age": 25
    }
    ```
    
    ```
    {
        "name": "Bob",
        "age": 31
    }
    ```
    
    ```python
    print("Hello, World!")
    ```
    """
    json_objects = extract_json(input_text)
    if json_objects:
        for obj in json_objects:
            print(obj)
    else:
        print("No valid JSON objects found.")
